import time
from .config import settings

class SnowflakeGenerator:
    def __init__(self, worker_id: int = settings.WORKER_ID, datacenter_id: int = settings.DATACENTER_ID):
        self.worker_id = worker_id
        self.datacenter_id = datacenter_id
        self.sequence = 0
        self.last_timestamp = -1

    def _get_timestamp(self) -> int:
        return int(time.time() * 1000)

    def _wait_next_millis(self, last_timestamp: int) -> int:
        timestamp = self._get_timestamp()
        while timestamp <= last_timestamp:
            timestamp = self._get_timestamp()
        return timestamp

    def generate(self) -> int:
        timestamp = self._get_timestamp()

        if timestamp < self.last_timestamp:
            raise Exception("Clock moved backwards!")

        if timestamp == self.last_timestamp:
            self.sequence = (self.sequence + 1) & 4095  # 12位序列号
            if self.sequence == 0:
                timestamp = self._wait_next_millis(self.last_timestamp)
                self.sequence = 0
        else:
            self.sequence = 0

        self.last_timestamp = timestamp

        # 组合ID
        # 1位符号位 + 41位时间戳 + 5位数据中心ID + 5位工作机器ID + 12位序列号
        id = ((timestamp - 1288834974657) << 22) | \
             (self.datacenter_id << 17) | \
             (self.worker_id << 12) | \
             self.sequence

        return id

# 创建全局ID生成器实例
snowflake = SnowflakeGenerator() 